import torch

class FScoreLoss(torch.nn.Module):

    def __init__(self, beta, eps=1e-7):
        super().__init__()
        self.beta = beta
        self.eps = eps

    def forward(self, y_true, y_pred, grad=True):
        tp = (y_true * y_pred).sum().to(torch.float32)
        fn = ((1 - y_true) * y_pred).sum().to(torch.float32)
        fp = (y_true * (1 - y_pred)).sum().to(torch.float32)
        precision = tp / (tp + fp + self.eps)
        recall = tp / (tp + fn + self.eps)

        f_score_loss = (1 + self.beta**2) * (precision * recall) / (
            (self.beta**2) * precision + recall + self.eps)

        # print(
        #     f'precision = {precision}, recall = {recall}, f_score = {f_score_loss}'
        # )

        return -1 * f_score_loss
